Are we initialized kv cache with different values?

This is for checking the numerical value with deterministic input for validation purpose. I can revert them back to zeros.

Are you copying over the quantization parameters from kv mode to prefill mode?

Yes, since the current bottleneck of calibration time is using prefill mode for generating user prompts.
I think the future scenario we want to try is to generate all the special tokens for target model and merge them with task calibration data. Have prefill model to iterate over it and use prefill's encoding for kv. We could save the user prompt generation with this approach.

what is the mechanism to make sure the encoding align correctly?

I'm worried about the accuracy too if we get rid of prefill calibration, do you think if we generate prompt + output using fp32 model (pre observers) as discussed in PR #17786 and run prefill + decode as before with skip_generate might yield better accuracy, rather than getting rid of entire prefill calibration?

prefill calibration ideally is not needed because decode see all the generated tokens too and prefill graph and decode graph should be the same. I remember @haowhsu-quic mentioned we insert the kv cache of the output of prefill and connect to the decode of the input to make sure those quant nodes are also calibrated. I did a comparison for quant params between prefill and decode in the past and they are very very close. I'm trying to figure out if this PR handle kv cache differently than before.

I guess my question is prefill "sees" previous tokens and as attention block will take these into consideration while generating kv-cache.

For weights, what you said makes sense as they are idempotent from the math standpoint, maybe we should just check the PPL on-device?

CC: @metascroy, @kimishpatel if you have any thoughts on this.

The mechanism is to compare the topology order of ops in both graph where each op in the same order should have identical nn_module_stack. The number / type of QDQ pairs in nodes' (call_function / placeholder) users are required to be identical as well.

Will have a fix for condition prefill_ar_len == max_seq_len, thanks for identifying this.

Ran comparison over the weekend, observed great speedups overall, calibration time is cut in half overall time is cut by ~33% for the qwen model i'm benchmarking. I will do a final PPL check on-device and we will be good to go. Thanks again for working on this.

I think we also should find ways to cut per-iteration decode timing.

┌─────────────────────┬──────────┬────────┬─────────────────────┐
│                     │ Baseline │   PR   │       Savings       │
├─────────────────────┼──────────┼────────┼─────────────────────┤
│ DECODE calibration  │ 2h 59m   │ 2h 49m │ ~10m (noise)        │
├─────────────────────┼──────────┼────────┼─────────────────────┤
│ PREFILL calibration │ 2h 34m   │ 15.5s  │ 2h 34m              │
├─────────────────────┼──────────┼────────┼─────────────────────┤
│ Quantization total  │ 5h 33m   │ 2h 50m │ 2h 43m (49% faster) │
├─────────────────────┼──────────┼────────┼─────────────────────┤
│ Compile total       │ 2h 35m   │ 2h 35m │ ~0 (noise)          │
├─────────────────────┼──────────┼────────┼─────────────────────┤
│ End-to-end          │ 8h 13m   │ 5h 30m │ 2h 43m (33% faster) │
└─────────────────────┴──────────┴────────┴─────────────────────┘

@haowhsu-quic PPL looks fine.

Let's go with the change. Can you please rebase.

The mechanism is to compare the topology order of ops in both graph where each op in the same order should have identical nn_module_stack. The number / type of QDQ pairs in nodes' (call_function / placeholder) users are required to be identical as well.

Is there any particular reason that we remove the custom annotation? I understand that the custom annotation caused prefill/decode have different graphs, but how about the quantization params for kv cache?

I think quantization params for placeholders (kv cache) will also be shared between prefill & decode. They should be the same under identical calibration data. May I learn more about your concern?